from pyNastran.converters.nastran.nastran_to_vtk import nastran_to_vtk
import meshio
import vtk
import os

def Op2ToVtk(op2_path: str, vtk_dir: str):
    # 判断文件是否存在
    is_exist = os.path.exists(op2_path)
    if not is_exist:
        print("[ERROR] Op2(" + op2_path + ") is not exist.")

    vtk_path = ""
    try:
        # op2文件名
        op2_file_name = os.path.basename(op2_path)
        file_name, ext = os.path.splitext(op2_file_name)

        vtk_file_name = file_name + ".vtu"
        vtk_path = vtk_dir + "/" + vtk_file_name
        nastran_to_vtk(op2_path, op2_path, vtk_path)

        # 读取vtu文件
        reader = vtk.vtkXMLUnstructuredGridReader()
        reader.SetFileName(vtk_path)
        reader.Update()
        read_data = reader.GetOutput()
        field_data = read_data.GetFieldData()
        num_of_arrays = field_data.GetNumberOfArrays()
        for i in range(num_of_arrays):
            field_data.RemoveArray(0)
        writer = vtk.vtkXMLUnstructuredGridWriter()
        writer.SetFileName(vtk_path)
        writer.SetDataModeToAscii()
        writer.SetInputData(read_data)
        writer.Update()
    except:
        print("[ERROR] Op2(" + op2_path + ") to vtk(" + vtk_dir + ") failed.")

    return vtk_path

def VtkToGmsh(vtk_path: str, gmsh_dir: str):
    # 判断文件是否存在
    is_exist = os.path.exists(vtk_path)
    if not is_exist:
        print("[ERROR] Vtk(" + vtk_path + ") is not exist.")

    try:
        # vtk文件名
        vtk_file_name = os.path.basename(vtk_path)
        file_name, ext = os.path.splitext(vtk_file_name)

        gmsh_file_name = file_name + ".msh"
        gmsh_path = gmsh_dir + "/" + gmsh_file_name

        mesh = meshio.read(vtk_path)
        mesh_without_field = meshio.Mesh(
            points=mesh.points,
            cells=mesh.cells,
            # cell_data=mesh.cell_data,
            # point_data=mesh.point_data,
        )
        meshio.write(gmsh_path, mesh_without_field, file_format="gmsh22", binary=False)
    except:
        print("[ERROR] Vtk(" + vtk_path + ") to gmsh(" + gmsh_path + ") failed.")

def ConvertOp2(op2_path: str, out_dir: str):
    # 判断文件是否存在
    is_exist = os.path.exists(op2_path)
    if not is_exist:
        print("[ERROR] Op2(" + op2_path + ") is not exist.")

    try:
        vtk_path = Op2ToVtk(op2_path, out_dir)
        VtkToGmsh(vtk_path, out_dir)
    except:
        print("[ERROR] Vtk(" + vtk_path + ") to gmsh(" + gmsh_path + ") failed.")